/**
 * Haskolinn i Reykjavik
 * Date		: 06. February 2010
 * Subject  : Computer Networking
 * Project  : TFTP
 * Students : Ingvi, Kjartan, Ragnar
 */

import java.net.*;
import java.io.*;

public class tftp {

    public final static int DEFAULT_PORT = 69;
    public final static int MODE_NETASCII = 0;
    public final static int MODE_BINARY = 1;
    public final static int OP_RRQ = 1;
    public final static int OP_WRQ = 2;
    public final static byte OP_DATAPACKET = 3;
    public final static byte OP_ACK = 4;
    public final static byte OP_ERROR = 5;
    public final static int PACKET_SIZE = 516;

    public static void main(String[] args)
    {
        //Checks for correct user input couldn't have !=3 because of *.*
        if (args.length < 3) {
            System.err.println("Usage: tftp [IP_TALA] [PUT/GET] [FILENAME] ");
            System.exit(1);
        }

        String host = args[0];  //Sets the hostname to variable
        String command = args[1].toUpperCase(); //Sets the PUT/GET to UpperCase in case of LowerCase input from user
        String fileName = args[2];  //Sets the filename to a variable

        DatagramSocket socket = null;
        InetAddress tftp_server = null;
        byte[] request = null;
        byte[] buffer = null;
        DatagramPacket outBound = null;
        DatagramPacket inBound = null;
        ByteArrayOutputStream byteOut = new ByteArrayOutputStream();

        try {
            tftp_server = InetAddress.getByName(host);
            socket = new DatagramSocket();
        }
        catch (IOException ex) {
            System.err.println(ex);
            System.exit(1);
        }

        //If it is a PUT input from user that is to upload a file to server
        if (command.equals("PUT")) {

            if(args.length > 3)
            {
                //Creates a variable for the current working directory
                String curDir = System.getProperty("user.dir");
                File dir = new File(curDir);

                //Puts all file in current directory into an array
                String[] children = dir.list();

                //Uploads all the files in the current directory
                for(int i=0;i<children.length;i++)
                {
                    File file = new File(children[i]);
                    sendFile(fileName, file, socket, tftp_server, request, buffer, outBound, inBound, byteOut);
                }

                //DEBUGGING
                //for(int i=0;i<children.length;i++)
                //System.out.println(children[i]);
            }
            else
            {
                //Creates a file for selected filename
                File file = new File(fileName);
                sendFile(fileName, file, socket, tftp_server, request, buffer, outBound, inBound, byteOut);
            }
        }

        //If it is a GET input from user that is to download a file from server
        if (command.equals("GET")) {

            try {
                request = createData(OP_RRQ, fileName, MODE_BINARY).toByteArray();
                outBound = new DatagramPacket(request, request.length, tftp_server, DEFAULT_PORT);
                socket.send(outBound);

                do {
                    buffer = new byte[PACKET_SIZE];
                    inBound = new DatagramPacket(buffer, buffer.length, tftp_server, socket.getLocalPort());
                    socket.receive(inBound);

                    //Getting the first 4 characters from the bufffer
                    byte[] opCode = {buffer[0], buffer[1]};

                    // Check to see if we received a data packet
                    if (opCode[1] == OP_ERROR) {
                        String errorCode = new String(buffer, 3, 1);
                        String errorText = new String(buffer, 4, inBound.getLength() - 4);
                        System.err.println("Error: " + errorCode + " " + errorText);
                    }

                    if (opCode[1] == OP_DATAPACKET) {
                        // Check for the packets block number
                        byte[] UDPBlockNr = {buffer[2], buffer[3]};

                        DataOutputStream dos = new DataOutputStream(byteOut);
                        dos.write(inBound.getData(), 4, inBound.getLength() - 4);

                        byte[] ACK = {0, OP_ACK, UDPBlockNr[0], UDPBlockNr[1]};

                        DatagramPacket ack = new DatagramPacket(ACK, ACK.length, tftp_server, inBound.getPort());
                        socket.send(ack);
                    }
                } while (!isLastPacket(inBound));
                writeFile(byteOut, fileName);
            }
            catch (IOException ex) {
                System.err.println(ex);
            }
        }
    }

    private static void sendFile(String fileName, File file, DatagramSocket socket, InetAddress tftp_server,
                                 byte[] request, byte[] buffer, DatagramPacket outBound, DatagramPacket inBound,
                                 ByteArrayOutputStream byteOut)
    {
        try {
                FileInputStream input = new FileInputStream(file);

                byte fileContent[] = new byte[512];

                //Send WRQ
                request = createData(OP_WRQ, fileName, MODE_BINARY).toByteArray();
                outBound = new DatagramPacket(request, request.length, tftp_server, DEFAULT_PORT);
                socket.send(outBound);

                //Receive ACK from server
                buffer = new byte[4];
                inBound = new DatagramPacket(buffer, buffer.length, tftp_server, socket.getLocalPort());
                socket.receive(inBound);

                int blockNr = 1;

                //while(input.available() > 512 && !isLastPacket(inBound)){
                while(input.available() > 1){

                    input.read(fileContent,0,512);
                    System.out.println(input.available());
                    DataOutputStream dos = new DataOutputStream(byteOut);
                    dos.writeShort(3);
                    dos.writeShort(blockNr);
                    dos.write(fileContent);

                    blockNr = blockNr + 1;
                    outBound = new DatagramPacket(byteOut.toByteArray(),byteOut.toByteArray().length,
                                                            tftp_server,inBound.getPort());
                    socket.send(outBound);
                    byteOut = new ByteArrayOutputStream();
                    socket.receive(inBound);
                }

                socket.close();
                input.close();
        }
        catch (IOException ex) {
                System.err.println(ex);
        }
    }

    private static void writeFile(ByteArrayOutputStream toFile, String fileName) {

        try {
            FileOutputStream output;
            File file;
            file = new File(fileName);
            output = new FileOutputStream(file);
            output.write(toFile.toByteArray());
            output.close();
        }
        catch (FileNotFoundException ex) {
            System.out.println("FileNotFoundException: " + ex);
        }
        catch (IOException ioe) {
            System.out.println("IOException: " + ioe);
        }

    }

    private static boolean isLastPacket(DatagramPacket pakki) {
        if (pakki.getLength() < 512)
            return true;
        else
            return false;
    }

    //Create data for reguestType (RRQ/WRQ), Filename, Transfermode (ASCII/BINARY), 
    private static ByteArrayOutputStream createData(int requestType, String fileName, int transferMode) throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        DataOutputStream dos = new DataOutputStream(baos);
        dos.writeShort(requestType);
        dos.writeBytes(fileName);
        dos.writeBytes("\0");
        if (transferMode == MODE_NETASCII) {
            dos.writeBytes("netascii");
        } else if (transferMode == MODE_BINARY) {
            dos.writeBytes("octet");
        }
        dos.writeBytes("\0");
        return baos;
    }
}
